package com.lxw.station.chatgpt.qianfan;

import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpEntity;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;

import java.nio.file.Files;
import java.nio.file.Paths;

@Slf4j
public class HelloQianFan {

    // 替换为你的API Key和Secret Key
    private static final String API_KEY = "K6xC6oTaxSQwnVz8AOYpWj8U";
    private static final String SECRET_KEY = "yInF3mvC35LdzyECJtYT1FsJHTo2afmV";
    private static String accessToken = "24.407e4923f7ac367f85b2d88f167e00ac.2592000.1722477846.282335-86466266";
    private static  RestTemplate restTemplate = new RestTemplate();
    private static final String ROLE_USER = "user";
    private static final String ROLE_ASSISTANT = "assistant";

    public static void main(String[] args) {
        Logger logger = (Logger) LoggerFactory.getLogger("ROOT");
        logger.setLevel(Level.INFO);
        getAccessToken();
        generateTitleAndDesc();
//        chat();
    }


    @SneakyThrows
    private static String generateTitleAndDesc() {
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token="+accessToken;
        QianFanParam qianFanParam = new QianFanParam();
        // 问
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_USER).setContent("你好,我给你一段文本,你能帮我生成一段描述和起个标题吗?返回数据格式:纯文本格式"));
        HttpEntity<QianFanParam> request = new HttpEntity<>(qianFanParam);
        QianFanChatResponse body = restTemplate.postForEntity(url, request, QianFanChatResponse.class).getBody();
        System.out.println(body.getResult());
        // 添加聊天记录
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_ASSISTANT).setContent(body.getResult()));
        // 问
        String readMeContent = new String(Files.readAllBytes(Paths.get("src/main/resources/tmp/README.md")));
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_USER).setContent(readMeContent));
        body = restTemplate.postForEntity(url, request, QianFanChatResponse.class).getBody();
        System.out.println("=============");
        System.out.println(body.getResult());
        return null;
    }
    private static String chat() {
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token="+accessToken;
        QianFanParam qianFanParam = new QianFanParam();
        // 问
        String me = "你好";
        System.out.println("me : " + me);
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_USER).setContent(me));
        HttpEntity<QianFanParam> request = new HttpEntity<>(qianFanParam);
        QianFanChatResponse body = restTemplate.postForEntity(url, request, QianFanChatResponse.class).getBody();
        String ai = body.getResult();
        System.out.println("ai : " + ai);
        // 添加聊天记录
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_ASSISTANT).setContent(body.getResult()));
        // 问
        me = "我在秦皇岛,周末可以去哪里玩?";
        System.out.println("me : " + me);
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_USER).setContent(me));
        body = restTemplate.postForEntity(url, request, QianFanChatResponse.class).getBody();
        ai = body.getResult();
        System.out.println("ai : " + ai);
        // 添加聊天记录
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_ASSISTANT).setContent(body.getResult()));
        // 问
        me = "周末这里的天气怎么样";
        System.out.println("me : " + me);
        qianFanParam.addMessage(new QianFanParam.Message().setRole(ROLE_USER).setContent(me));
        body = restTemplate.postForEntity(url, request, QianFanChatResponse.class).getBody();
        ai = body.getResult();
        System.out.println("ai : " + ai);
        return null;
    }

    private static void getAccessToken() {
        if (accessToken == null) {
            String url = "https://aip.baidubce.com/oauth/2.0/token?client_id=${apiKey}&client_secret=${apiSecret}&grant_type=client_credentials";
            url = url.replace("${apiKey}", API_KEY).replace("${apiSecret}", SECRET_KEY);
            ResponseEntity<QianFanAuthToken> qianFanAuthToken = restTemplate.getForEntity(url, QianFanAuthToken.class);
            System.out.println("qianFanAuthToken = " + qianFanAuthToken);
            if (qianFanAuthToken.getBody().getError_description() != null) {
                log.error("获取千帆大模型token请求失败:{}", qianFanAuthToken.getBody().getError_description());
            }
            accessToken = qianFanAuthToken.getBody().getAccess_token();
            System.out.println(accessToken);
        }
    }
}
